import os
import time
import random
import json
from tqdm import tqdm  # 进度条工具（当前代码未直接使用，预留接口）

import torch
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_  # 梯度裁剪工具
from torch.utils.data import DataLoader  # 数据加载器
from tensorboardX import SummaryWriter  # TensorBoard可视化工具
from torch.optim.lr_scheduler import StepLR, MultiStepLR  # 学习率调度器

import numpy as np
from configs.opts import parser  # 命令行参数解析器
from model.main_model import weak_main_model as main_model  # 弱监督主模型
from utils import AverageMeter, Prepare_logger, get_and_save_args  # 工具函数/类
from utils.Recorder import Recorder  # 实验记录器（代码备份、日志保存）
from dataset.AVE_dataset_weak import AVEDataset  # 弱监督专用AVE数据集类

# =================================  随机种子配置 ============================
# 固定随机种子以确保实验可复现（弱监督版本种子设为666，与有监督版本区分）
SEED = 666
random.seed(SEED)
np.random.seed(seed=SEED)
torch.manual_seed(seed=SEED)
torch.cuda.manual_seed(seed=SEED)
torch.backends.cudnn.deterministic = True  # 禁用cudnn非确定性算法
torch.backends.cudnn.benchmark = False  # 禁用benchmark模式
# =============================================================================

# 加载弱监督模型配置文件（JSON格式，与有监督版本的配置文件不同）
config_path = 'configs/weak.json'
with open(config_path) as fp:
    config = json.load(fp)
print(config)  # 打印配置信息，确认加载正确


def main():
    # 全局变量声明（跨函数共享关键参数和统计信息）
    global args, logger, writer, dataset_configs
    global best_accuracy, best_accuracy_epoch
    best_accuracy, best_accuracy_epoch = 0, 0  # 记录最佳准确率及对应epoch

    # 加载并合并配置参数（命令行参数覆盖默认配置）
    dataset_configs = get_and_save_args(parser)
    parser.set_defaults(**dataset_configs)
    args = parser.parse_args()

    # GPU设备配置（指定使用的GPU编号）
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    '''创建实验快照目录：用于备份代码和保存模型'''
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)
    # 若从checkpoint恢复训练，使用checkpoint所在目录作为快照目录
    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    # 初始化日志记录器（区分训练和评估模式）
    logger = Prepare_logger(args, eval=args.evaluate)

    # 打印实验配置（训练模式下）
    if not args.evaluate:
        logger.info(f'\nCreating folder: {args.snapshot_pref}')
        logger.info('\nRuntime args\n\n{}\n'.format(json.dumps(vars(args), indent=4)))
    else:
        logger.info(f'\nLog file will be save in {args.snapshot_pref}/Eval.log.')

    '''数据集加载：初始化弱监督训练和测试数据加载器'''
    train_dataloader = DataLoader(
        AVEDataset('./data/', split='train'),  # 弱监督训练集（含干净样本+噪声样本）
        batch_size=args.batch_size,  # 训练批次大小
        shuffle=True,  # 训练集打乱顺序
        num_workers=8,  # 数据加载线程数
        pin_memory=True  # 锁定内存，加速GPU数据传输
    )

    test_dataloader = DataLoader(
        AVEDataset('./data/', split='test'),  # 测试集（与有监督版本一致）
        batch_size=args.test_batch_size,  # 测试批次大小
        shuffle=False,  # 测试集不打乱
        num_workers=8,
        pin_memory=True
    )

    '''模型配置：初始化弱监督模型、优化器、调度器和损失函数'''
    mainModel = main_model(config["model"])  # 初始化弱监督主模型
    mainModel = nn.DataParallel(mainModel).cuda()  # 多GPU并行训练（移至GPU）
    learned_parameters = mainModel.parameters()  # 待优化的模型参数
    optimizer = torch.optim.Adam(learned_parameters, lr=args.lr)  # Adam优化器
    # 学习率调度器：多步衰减（弱监督版本的衰减节点为[10,20,40]，与有监督版本不同）
    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 40], gamma=0.5)
    # 损失函数：适配弱监督多标签场景
    criterion = nn.BCEWithLogitsLoss().cuda()  # 二分类交叉熵（备用）
    criterion_event = nn.MultiLabelSoftMarginLoss().cuda()  # 多标签软边际损失（核心损失）

    '''从checkpoint恢复训练/评估'''
    if os.path.isfile(args.resume):
        logger.info(f"\nLoading Checkpoint: {args.resume}\n")
        mainModel.load_state_dict(torch.load(args.resume))  # 加载模型权重
    elif args.resume != "" and (not os.path.isfile(args.resume)):
        raise FileNotFoundError  # 若指定resume但文件不存在，抛出异常

    '''仅执行评估（不训练）'''
    if args.evaluate:
        logger.info(f"\nStart Evaluation..")
        validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch=0, eval_only=True)
        return

    '''TensorBoard可视化和代码备份初始化'''
    writer = SummaryWriter(args.snapshot_pref)  # TensorBoard日志写入器
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")  # 实验记录器
    recorder.writeopt(args)  # 保存实验参数到文件

    '''训练与测试主循环'''
    for epoch in range(args.n_epoch):
        # 训练一个epoch，返回平均损失
        loss = train_epoch(mainModel, train_dataloader, criterion, criterion_event, optimizer, epoch)

        # 按指定频率评估模型（或最后一个epoch强制评估）
        if ((epoch + 1) % args.eval_freq == 0) or (epoch == args.n_epoch - 1):
            acc = validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch)
            # 若当前准确率优于历史最佳，保存模型checkpoint
            if acc > best_accuracy:
                best_accuracy = acc
                best_accuracy_epoch = epoch
                save_checkpoint(
                    mainModel.state_dict(),
                    top1=best_accuracy,
                    task='WeaklySupervised',  # 标记为弱监督任务，区分模型文件
                    epoch=epoch + 1,
                )
            print("-----------------------------")
            print("best acc and epoch:", best_accuracy, best_accuracy_epoch)
            print("-----------------------------")
        # 更新学习率
        scheduler.step()


def train_epoch(model, train_dataloader, criterion, criterion_event, optimizer, epoch):
    """
    弱监督单轮训练函数：适配弱监督标签特性，简化损失计算（仅保留类别损失）
    Args:
        model: 弱监督模型
        train_dataloader: 弱监督训练数据加载器
        criterion: 二分类损失函数（备用，未使用）
        criterion_event: 多标签损失函数（核心）
        optimizer: 优化器
        epoch: 当前epoch编号
    Returns:
        本轮训练的平均损失
    """
    # 初始化统计指标计数器
    batch_time = AverageMeter()  # 批次处理时间
    data_time = AverageMeter()   # 数据加载时间
    losses = AverageMeter()      # 损失值
    train_acc = AverageMeter()   # 训练准确率（弱监督训练中未实际计算，设为0）
    end_time = time.time()       # 计时起点

    model.train()  # 设置模型为训练模式
    model.double()  # 转换模型为double精度（适配输入特征类型）
    optimizer.zero_grad()  # 初始化梯度为0

    # 遍历训练集批次
    for n_iter, batch_data in enumerate(train_dataloader):
        # 更新数据加载时间
        data_time.update(time.time() - end_time)

        '''加载批次数据并预处理'''
        visual_feature, audio_feature, labels = batch_data  # 视觉特征、音频特征、弱监督标签
        labels = labels.double().cuda()  # 标签转换为double精度并移至GPU

        '''模型前向传播（弱监督模型输出与有监督不同）'''
        is_event_scores, raw_logits, event_scores = model(visual_feature, audio_feature)
        # 输出说明：is_event_scores（事件置信度）、raw_logits（原始类别分数）、event_scores（归一化类别分数）

        '''损失计算：弱监督仅使用多标签类别损失'''
        loss_event_class = criterion_event(event_scores, labels)  # 多标签软边际损失
        loss = loss_event_class  # 总损失（无其他辅助损失，与有监督多损失融合不同）
        loss.backward()  # 反向传播计算梯度

        '''计算训练准确率（弱监督训练中暂不计算，返回0占位）'''
        acc = torch.tensor([0])
        train_acc.update(acc.item(), visual_feature.size(0) * 10)  # 按样本数加权更新

        '''梯度裁剪：防止梯度爆炸'''
        if args.clip_gradient is not None:
            total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)

        '''优化器更新模型参数'''
        optimizer.step()
        optimizer.zero_grad()  # 重置梯度

        '''更新统计指标'''
        losses.update(loss.item(), visual_feature.size(0) * 10)  # 按时序步长加权损失
        batch_time.update(time.time() - end_time)
        end_time = time.time()

        '''记录TensorBoard日志（迭代级）'''
        writer.add_scalar('Train_data/loss', losses.val, epoch * len(train_dataloader) + n_iter + 1)

        '''记录TensorBoard日志（epoch级）'''
        writer.add_scalar('Train_epoch_data/epoch_loss', losses.avg, epoch)

    return losses.avg


@torch.no_grad()  # 禁用梯度计算，加速评估并节省内存
def validate_epoch(model, test_dataloader, criterion, criterion_event, epoch, eval_only=False):
    """
    弱监督单轮评估函数：使用测试集计算模型准确率，适配弱监督输出格式
    Args:
        model: 待评估的弱监督模型
        test_dataloader: 测试数据加载器
        criterion: 二分类损失函数（未使用）
        criterion_event: 多标签损失函数（未使用）
        epoch: 当前epoch编号
        eval_only: 是否仅执行评估（不记录TensorBoard）
    Returns:
        测试集平均准确率
    """
    # 初始化统计指标计数器
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()  # 弱监督评估中未计算损失，仅占位
    accuracy = AverageMeter()
    end_time = time.time()

    model.eval()  # 设置模型为评估模式
    model.double()  # 转换为double精度

    # 遍历测试集批次
    for n_iter, batch_data in enumerate(test_dataloader):
        data_time.update(time.time() - end_time)

        '''加载批次数据并预处理'''
        visual_feature, audio_feature, labels = batch_data
        labels = labels.cuda()
        bs = visual_feature.size(0)  # 批次大小

        '''模型前向传播'''
        is_event_scores, raw_logits, event_scores = model(visual_feature, audio_feature)

        '''计算评估准确率（使用弱监督适配的准确率计算逻辑）'''
        acc = compute_accuracy_supervised(is_event_scores, raw_logits, labels)
        accuracy.update(acc.item(), bs)  # 按批次大小加权更新

        '''更新统计指标'''
        batch_time.update(time.time() - end_time)
        end_time = time.time()

    '''记录TensorBoard日志（非仅评估模式）'''
    if not eval_only:
        writer.add_scalar('Val_epoch/Accuracy', accuracy.avg, epoch)

    # 打印本轮评估的准确率
    logger.info(
        f'**************************************************************************\t'
        f"\tEvaluation results (acc): {accuracy.avg:.4f}%."
    )
    return accuracy.avg


def compute_accuracy_supervised(is_event_scores, event_scores, labels):
    """
    适配弱监督的准确率计算函数：区分事件与背景，匹配弱监督模型输出格式
    Args:
        is_event_scores: 事件置信度分数 (batch_size, seq_len)
        event_scores: 原始类别分数 (batch_size, num_classes)
        labels: 测试集真实标签 (batch_size, seq_len, num_classes+1)
    Returns:
        准确率（百分比）
    """
    _, targets = labels.max(-1)  # 提取真实标签的类别索引（包含背景类）
    # 事件置信度sigmoid激活（转换为0-1概率）
    is_event_scores = is_event_scores.sigmoid()
    scores_pos_ind = is_event_scores > 0.5  # 置信度>0.5视为事件
    scores_mask = scores_pos_ind == 0  # 背景掩码（置信度<=0.5）
    # 事件类别预测（取分数最高的类别）
    _, event_class = event_scores.max(-1)
    # 构建预测结果：事件位置填充类别，背景位置填充背景索引（28）
    pred = scores_pos_ind.long()
    pred *= event_class[:, None]
    pred[scores_mask] = 28  # 28表示背景类
    # 计算正确预测数并求准确率
    correct = pred.eq(targets)
    correct_num = correct.sum().double()
    acc = correct_num * (100. / correct.numel())  # 转换为百分比
    return acc


def compute_accuracy_weak(event_scores, labels):
    """
    弱监督专用准确率计算函数（备用）：直接对类别分数取最大值，适配多标签场景
    Args:
        event_scores: 类别分数 (batch_size, seq_len, num_classes)
        labels: 真实标签 (batch_size, seq_len, num_classes+1)
    Returns:
        准确率（百分比）
    """
    _, pred = event_scores.max(-1)  # 预测类别（取最高分）
    pred = pred.transpose(1, 0).contiguous()  # 调整维度
    _, target = labels.max(-1)  # 真实类别
    correct = pred.eq(target)  # 对比预测与真实标签
    correct_num = correct.sum().double()
    acc = correct_num * (100. / correct.numel())
    return acc


def save_checkpoint(state_dict, top1, task, epoch):
    """
    保存模型checkpoint（弱监督版本标记任务类型为"WeaklySupervised"）
    Args:
        state_dict: 模型状态字典
        top1: 当前最佳准确率
        task: 任务类型（弱监督）
        epoch: 当前epoch编号
    """
    # 构建checkpoint文件名（含epoch、准确率、任务类型）
    model_name = f'{args.snapshot_pref}/model_epoch_{epoch}_top1_{top1:.3f}_task_{task}_best_model.pth.tar'
    torch.save(state_dict, model_name)  # 保存模型权重


if __name__ == '__main__':
    main()  # 启动弱监督训练/评估流程